import os

from easydict import EasyDict
from pyecharts import options as opts
from pyecharts.charts import Sunburst

from investment.constants import ASSET_NAME_DICT


class AssetRecord:
    def __init__(self, cls, subclass, name, money):
        self.cls = cls.name
        self.subclass = subclass.name
        self.name = name
        self.money = money
        self.ratio = None

    def set_ratio(self, sum_money):
        self.ratio = self.money / sum_money


class AssetTable:
    def __init__(self, records_list):
        assert isinstance(records_list, list)
        self.records = records_list[:]
        self.sum_money = 0
        for rd in self.records:
            self.sum_money += rd.money
        for rd in self.records:
            rd.set_ratio(self.sum_money)
        self.cls_ratio = EasyDict()

        # get big class ratio
        for rd in self.records:
            assert rd.ratio is not None
            if rd.cls not in self.cls_ratio.keys():
                self.cls_ratio[rd.cls] = {}
                self.cls_ratio[rd.cls].ratio = rd.ratio
                # new subclass
                self.cls_ratio[rd.cls].subclass = {rd.subclass: {'records': [{rd.name: rd.ratio}], 'ratio': rd.ratio}}
            else:
                self.cls_ratio[rd.cls].ratio += rd.ratio
                if rd.subclass not in self.cls_ratio[rd.cls].subclass.keys():
                    # new subclass
                    self.cls_ratio[rd.cls].subclass[rd.subclass] = {'records': [{rd.name: rd.ratio}], 'ratio': rd.ratio}
                else:
                    # add to old subclass
                    self.cls_ratio[rd.cls].subclass[rd.subclass]['records'].append({rd.name: rd.ratio})
                    self.cls_ratio[rd.cls].subclass[rd.subclass]['ratio'] += rd.ratio

    def save_to_txt(self, f_path):
        with open(f_path, 'w') as f:
            for rd in self.records:
                f.write(",".join([rd.cls, rd.subclass, rd.name, str(rd.money)]))
                f.write("\n")


def generate_pie_charts(assets_table: AssetTable, out_name: str = "投资资产配置"):
    assert isinstance(assets_table, AssetTable)
    data = []
    for ky in assets_table.cls_ratio.keys():
        sub_cls_list = []
        for sub_ky in assets_table.cls_ratio[ky].subclass.keys():
            record_itm_list = []
            for record in assets_table.cls_ratio[ky].subclass[sub_ky].records:
                record_name = list(record.keys())[0]+" ({:.1%})".format(list(record.values())[0])
                record_itm = opts.SunburstItem(name=record_name, value=list(record.values())[0])
                record_itm_list.append(record_itm)
            sub_cls = opts.SunburstItem(name=sub_ky, children=record_itm_list,
                                        value=assets_table.cls_ratio[ky].subclass[sub_ky].ratio)
            sub_cls_list.append(sub_cls)
        cls_name = ky + "({:.1%})".format(assets_table.cls_ratio[ky].ratio)
        cls_item = opts.SunburstItem(name=cls_name, children=sub_cls_list, value=assets_table.cls_ratio[ky].ratio)
        data.append(cls_item)

    sunburst = (
        Sunburst(init_opts=opts.InitOpts(width="900px", height="600px"))
        .add(series_name="", data_pair=data, radius=[0, "90%"])
        .set_global_opts(title_opts=opts.TitleOpts(title=out_name))
        .set_series_opts(label_opts=opts.LabelOpts(formatter="{b}"))
        .render(os.path.join("plots", out_name+".html"))
    )


def read_asset_table_from_txt(f_path):
    record_lists = []
    with open(f_path, 'r') as f:
        lines = f.readlines()
        # lines = lines.strip()
        for line in lines:
            items = line.strip().split(",")
            asset_record = AssetRecord(ASSET_NAME_DICT[items[0]], ASSET_NAME_DICT[items[1]], items[2], eval(items[3]))
            record_lists.append(asset_record)
    return AssetTable(records_list=record_lists)


if __name__ == "__main__":
    date = "20220616"
    root = os.path.join(os.getcwd(), r".\plots")
    asset_table = read_asset_table_from_txt(os.path.join(root, "assets_{}.txt".format(date)))
    generate_pie_charts(asset_table, f"投资资产配置_{date}")
